from typing import Any, Optional, Dict
from starlette.background import BackgroundTask
from starlette.responses import JSONResponse
from fastapi.encoders import jsonable_encoder


class FormatJSONResponse(JSONResponse):
    def __init__(self,
                 data: Any = None,
                 code: int = 0,
                 msg: str = 'success',
                 status_code: int = 200,
                 headers: Optional[Dict[str, str]] = None,
                 background: Optional[BackgroundTask] = None
                 ) -> None:
        content = jsonable_encoder({
            'code': code,
            'msg': msg,
            'data': data
        })

        status_code = int(status_code)
        super().__init__(content=content, status_code=status_code,
                         headers=headers, media_type='application/json', background=background)
